import os
from random import random, randint
from mlflow import log_metric, log_param, log_artifacts, create_experiment, get_experiment, \
                    start_run, delete_run, get_run, set_tag

'''
Purpose of this example is to test the MLFlow Tracking API by
- creating and getting experiment
- creating and deleting run tasks
- Loggings, including log_metric, log_param, log_artifacts
Furthermore, for the Metric data recorded, we will finally visualize it in MLFlow UI platform.
'''

if __name__ == "__main__":
    print("Running mlflow_tracking.py")

    # Create an experiment name, which must be unique and case sensitive
    # If you run it again with same experiment name, you will get 
    # error: mlflow.exceptions.MlflowException: Experiment 'mlflow tracking example' already exists.
    experiment_id = create_experiment("mlflowTrackingExample3")
    experiment = get_experiment(experiment_id)
    print("Name: {}".format(experiment.name))
    print("Experiment_id: {}".format(experiment.experiment_id))
    print("Artifact Location: {}".format(experiment.artifact_location))
    print("Tags: {}".format(experiment.tags))
    print("Lifecycle_stage: {}".format(experiment.lifecycle_stage))

    # 启动一个新的MLflow run task
    with start_run(experiment_id=experiment_id) as parent_run:
        set_tag("test.release", "1.1.0-RC")     # 添加一个tag
        log_param("parent", "yes")              # 第一种Log形式：Log parameters
        log_param("param1", randint(0, 100))    

        for epoch in range(0, 20):              # 第二种Log形式：Log metric
            log_metric(key="quality", value=2*epoch, step=epoch)

        if not os.path.exists("outputs"):
            os.makedirs("outputs")
        with open("outputs/test.txt", "w") as f:
            f.write("hello world!")
        log_artifacts("outputs")                # 第三种Log形式：Log artifacts

    run_id = parent_run.info.run_id
    #delete_run(run_id)     # If you want to visualize the matric results, do not delete this run.

    print("run_id: {}; lifecycle_stage: {}".format(run_id, get_run(run_id).info.lifecycle_stage))